# import paddlex as pdx
import paddle

from Deploy import deploy
from ModelTrain import model_train
from Seq2Seq import Seq2seq

dic_kw_path = "dataset/字向量kw.npy"
dic_wk_path = "dataset/字向量wk.npy"
seqmodel = Seq2seq(hidden_dim=256, num_layers=2, dropout=0.2)
try:
    layer_state_dict = paddle.load("work/model/final.pdparams")
    seqmodel.set_state_dict(layer_state_dict)
    optim = paddle.optimizer.Adam(parameters=seqmodel.parameters())
    optim.set_state_dict(paddle.load("work/model/final.pdopt"))
    print("加载成功")
except:
    optim = paddle.optimizer.Adam(parameters=seqmodel.parameters())
    print("加载失败")


print("请选择相关功能")
print("1.训练  2.部署")
while True:
    choose=int(input("按数字键选择："))
    if choose==1:
       mt=model_train()
       mt.basetrain()
    elif choose==2:
        dp = deploy(seqmodel, dic_kw_path, dic_wk_path)
        print("已经开启部署")
        dp.startDeploy()
    else:
        print("请重新输入")


# 以下为超参数设置

#
# model.evaluate(test_dataset)
# batch_id = 0
# for data in range(100):
#     datastr="你是哪位"#input("我说：")
#     data,d,db=train_dataset.get_tensor(datastr,100)
#     data=paddle.to_tensor(data)
#     d = paddle.to_tensor(d)
#     db = paddle.to_tensor(db)
#     predicts = seqmodel.infer(data,d,db)
#     down = test_dataset.get_word(predicts)
#     down_str = ""
#     for c in down[0]:
#         if c=="eos":
#             break
#         down_str += c
#     print("机器人说：", down_str)
#     batch_id += 1
#     if batch_id > 100:
#         break
